__all__ = ['parse_command_line', 'log_config', 'init', 'schedule_tasks', 'clean']


import argparse
from Crypto.Cipher import AES
import itertools
import json
import logging
from multiprocessing.pool import ThreadPool
import os
import re
import shutil
import ssl
import socket
import subprocess
import sys
import types
import urllib.request

import util


__author__  = 'ZL <837806295@qq.com>'
__product__ = 'fetchm3u8'
__version__ = '4.1.0'
__date__    = '2023/04/08'


log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())  # 如果宿主未配置日志处理器，则丢弃日志

args = types.SimpleNamespace()

password = None

m3u8_index = 'index.m3u8'
m3u8_extension = '.ts'

max_workers = 64


def is_sequence_line(str_):
    """检查是否多媒体序列号行，返回序列号部分
    """

    if not isinstance(str_, str) or not str_:
        log.debug('is sequence line ... no')
        return ''

    m = re.search(r'^\s*\#\s*EXT\-X\-MEDIA\-SEQUENCE\s*\:\s*([0-9]+)\s*?',
                  str_,
                  re.I)
    if m and m.group() and m.group(1):
        sequence = m.group(1)
        log.debug('is sequence line ... yes')
        return sequence

    log.debug('is sequence line ... no')
    return ''

def is_key_line(str_):
    """检查是否加密信息行
    """

    if not isinstance(str_, str) or not str_:
        log.debug('is key line ... no')
        return ''

    m = re.search(r'^\s*\#\s*EXT\-X\-KEY\s*\:', str_, re.I)
    if m and m.group():
        log.debug('is key line ... yes')
        return True

    log.debug('is key line ... no')
    return False

def is_ts_line(str_):
    """检查是否ts地址行

    如果除去空白符以“.ts”结尾，认为是ts地址行，不满足的话，如果除去空白符不以“#”开头的非空行，
    认为是ts地址行
    """

    if not isinstance(str_, str) or not str_:
        log.debug('is ts line ... no')
        return ''

    m = re.search(r'^\s*([^\s].*\.ts)\s*$', str_, re.I)
    if m and m.group() and m.group(1):
        ts_url = m.group(1)
        log.debug('is ts line ... yes')
        return ts_url

    m = re.search(r'^\s*([^#\s].*[^\s])\s*$', str_)
    if m and m.group() and m.group(1):
        ts_url = m.group(1)
        log.debug('is ts line ... yes')
        return ts_url

    m = re.search(r'^\s*([^#\s])\s*$', str_)
    if m and m.group() and m.group(1):
        ts_url = m.group(1)
        log.debug('is ts line ... yes')
        return ts_url

    log.debug('is ts line ... no')
    return ''

def get_encryption_method(str_):
    """从“#EXT-X-KEY”行中获取使用的加密方法

    从“#EXT-X-KEY”行中获取“METHOD=”后面指定的加密方式，如：
        输入：#EXT-X-KEY:METHOD=AES-128,URI="https://abc.com/key.key"
        返回：AES-128

        输入：#EXT-X-KEY:URI="https://abc.com/key.key",METHOD=AES-128
        返回：AES-128
    """

    if not isinstance(str_, str) or not str_:
        log.error('get encryption method ... failed | illegal input')
        return ''

    m = re.search(r'^\s*\#\s*EXT\-X\-KEY.+?METHOD\s*\='
                  r'\s*([a-zA-Z0-9_\-]+?)\s*\,',
                  str_,
                  re.I)
    if m and m.group() and m.group(1):
        method = m.group(1)
        log.debug('get encryption method ... successed')
        return method

    m = re.search(r'^\s*\#\s*EXT\-X\-KEY.+?METHOD\s*\='
                  r'\s*([a-zA-Z0-9_\-]+?)\s*$',
                  str_,
                  re.I)
    if m and m.group() and m.group(1):
        method = m.group(1)
        log.debug('get encryption method ... successed')
        return method

    log.error('get encryption method ... failed')
    return ''

def get_key_url(str_):
    """从“#EXT-X-KEY”行中获取密钥URL

    从“#EXT-X-KEY”行中获取“URI=”后面指定的密钥URL，如：
        输入：#EXT-X-KEY:METHOD=AES-128,URI="https://abc.com/key.key"
        返回：https://abc.com/key.key

        输入：#EXT-X-KEY:URI="https://abc.com/key.key",METHOD=AES-128
        返回：https://abc.com/key.key
    """

    if not isinstance(str_, str) or not str_:
        log.error('get key url ... failed | illegal input')
        return ''

    m = re.search(
        r'^\s*\#\s*EXT\-X\-KEY.+?URI\s*\=\s*\"*([^"\s,].*?[^"\s])\"*\s*\,',
        str_,
        re.I)
    if m and m.group() and m.group(1):
        url = m.group(1)
        log.debug('get key url ... successed')
        return url

    m = re.search(
        r'^\s*\#\s*EXT\-X\-KEY.+?URI\s*\=\s*\"*([^"\s,].*?[^"\s])\"*\s*$',
        str_,
        re.I)
    if m and m.group() and m.group(1):
        url = m.group(1)
        log.debug('get key url ... successed')
        return url

    m = re.search(r'^\s*\#\s*EXT\-X\-KEY.+?URI\s*\=\s*\"*([^"\s,])\"*\s*\,',
                  str_,
                  re.I)
    if m and m.group() and m.group(1):
        url = m.group(1)
        log.debug('get key url ... successed')
        return url

    m = re.search(r'^\s*\#\s*EXT\-X\-KEY.+?URI\s*\=\s*\"*([^"\s,])\"*\s*$',
                  str_,
                  re.I)
    if m and m.group() and m.group(1):
        url = m.group(1)
        log.debug('get key url ... successed')
        return url

    log.debug('get key url ... failed')
    return ''

def get_encryption_iv(str_):
    """从“#EXT-X-KEY”行中获取加密向量

    从“#EXT-X-KEY”行中获取“IV=”后面指定的加密向量，如：
        输入：#EXT-X-KEY:METHOD=AES-128,URI="https://abc.com/key.key",IV=0x000102030405060708090a0b0c0d0e0f
        返回：000102030405060708090a0b0c0d0e0f

        输入：#EXT-X-KEY:URI="https://abc.com/key.key",IV=0x000102030405060708090a0b0c0d0e0f,METHOD=AES-128
        返回：000102030405060708090a0b0c0d0e0f
    """

    if not isinstance(str_, str) or not str_:
        log.error('get encryption iv ... failed | illegal input')
        return ''

    m = re.search(r'^\s*\#\s*EXT\-X\-KEY.+?IV\s*\=\s*0x([0-9a-fA-F]{32})\s*\,',
                  str_,
                  re.I)
    if m and m.group() and m.group(1):
        iv = m.group(1)
        log.debug('get encryption iv ... successed')
        return iv

    m = re.search(r'^\s*\#\s*EXT\-X\-KEY.+?IV\s*\=\s*0x([0-9a-fA-F]{32})\s*$',
                  str_,
                  re.I)
    if m and m.group() and m.group(1):
        iv = m.group(1)
        log.debug('get encryption iv ... successed')
        return iv

    log.debug('get encryption iv ... failed')
    return ''

def parse_key_line(str_, url_root, relative_url_head):
    """分析加密信息行，提取加密信息

    以如下格式把key信息放入list并返回

    [
        encryption_method,
        key_url,
        encryption_iv,
    ]

    encryption_method: 加密方法
    key_url: 加密密钥的完整地址
    encryption_iv: 加密向量
    """

    method = get_encryption_method(str_)
    if not method:
        log.error('parse key line ... failed')
        return []

    if method.upper() == 'NONE':
        url = ''
        iv = ''
    else:
        url = util.complete_url(get_key_url(str_), url_root, relative_url_head)
        iv = get_encryption_iv(str_)

    encryption_info = [
        method,
        url,
        iv,
    ]

    log.debug('parse key line ... successed')
    return encryption_info

def parse_ts_line(ts_url, url_root, relative_url_head, sequence,
                  encryption_info_index):
    """分析ts地址行，提取ts信息

    以如下格式把ts信息放入list并返回

    [
        url,
        sequence,
        encrypt_info_index,
    ]

    url: 完整的地址
    sequence: 多媒体序列号
    encrypt_info_index: 关联到加密信息list的下标，指示所使用加密信息位置
    """

    ts_info = [
        util.complete_url(ts_url, url_root, relative_url_head),
        sequence,
        encryption_info_index,
    ]

    log.debug('parse ts line ... done')
    return ts_info

def get_task_root_dir(location, title):
    """获取单个任务存储路径，路径不存在则创建
    """

    task_root_dir = os.path.join(location, title)
    if os.path.exists(task_root_dir):
        if not os.path.isdir(task_root_dir):
            try:
                os.remove(task_root_dir)
                os.mkdir(task_root_dir)
            except Exception:
                log.exception('get task root dir ... exception ... exit')
                sys.exit(1)
    else:
        try:
            os.mkdir(task_root_dir)
        except Exception:
            log.exception('get task root dir ... exception ... exit')
            sys.exit(1)

    log.debug('get task root dir ... successed')
    return task_root_dir

def get_m3u8_dir(location, title, dir='m3u8'):
    """获取存放m3u8数据的目录，路径不存在则创建
    """

    task_root_dir = get_task_root_dir(location, title)
    m3u8_dir = os.path.join(task_root_dir, dir)
    if os.path.exists(m3u8_dir):
        if not os.path.isdir(m3u8_dir):
            try:
                os.remove(m3u8_dir)
                os.mkdir(m3u8_dir)
            except Exception:
                log.exception('get m3u8 dir ... exception ... exit')
                sys.exit(1)
    else:
        try:
            os.mkdir(m3u8_dir)
        except Exception:
            log.exception('get m3u8 dir ... exception ... exit')
            sys.exit(1)

    log.debug('get m3u8 dir ... successed')
    return m3u8_dir

def has_task(location, title, filename='task.cfg'):
    """确定单个任务基本配置文件是否存在
    """

    task_path = os.path.join(location, title, filename)
    if os.path.isfile(task_path):
        log.debug('has task ... yes')
        return True
    else:
        log.debug('has task ... no | not file: %s', task_path)
        return False

def check_task_content(url, lines, keys, tss):
    """检查基本任务配置内容是否有效

    如果行数小于6，则认为是无效的，如果没有ts行，则认为是无效的，如果lines与tss长度不满足特定
    要求，则认为是无效的，其它情况认为有效
    """

    if not isinstance(url, str):
        log.error('check task content ... error | url')
        return False

    if not isinstance(lines, list) or len(lines) < 6:
        log.error('check task content ... error | lines')
        return False

    if not isinstance(keys, list):
        log.error('check task content ... error | keys')
        return False

    if not isinstance(tss, list) or not tss:
        log.error('check task content ... error | tss')
        return False

    if len(lines) < 2*len(tss) + 3:
        log.error('check task content ... error | lines < 2*tss + 3')
        return False

    log.debug('check task content ... ok')
    return True

def check_task(task):
    """对基本任务配置文件做格式校验
    """

    if not isinstance(task, dict) or not task:
        log.error('check task ... error')
        return False

    keys = {'url', 'lines', 'keys', 'tss'}
    if task.keys() ^ keys:
        log.error('check task ... error')
        return False

    if not check_task_content(task['url'], task['lines'], task['keys'],
                              task['tss']):
        log.error('check task ... error')
        return False

    log.debug('check task ... ok')
    return True

def load_task(location, title, filename='task.cfg'):
    """加载基本任务配置文件

    加载基本任务配置文件，并做格式校验
    """

    task_root_dir = get_task_root_dir(location, title)
    task_path = os.path.join(task_root_dir, filename)
    if not os.path.isfile(task_path):
        log.error('load task ... failed | not file: %s', task_path)

    task = {}
    try:
        with open(task_path, 'r', encoding='utf-8') as f:
            task = json.load(f)
    except Exception:
        log.exception('load task ... exception')
        return {}

    if not check_task(task):
        log.error('load task ... failed')
        return {}

    log.debug('load task ... successed')
    return task

def save_task(task, location, title, filename='task.cfg'):
    """保存基本任务配置到文件
    """

    if not check_task(task):
        log.error('save task ... failed')
        return False

    task_root_dir = get_task_root_dir(location, title)
    task_path = os.path.join(task_root_dir, filename)
    try:
        with open(task_path, 'w', encoding='utf-8') as f:
            json.dump(task, f, indent=2)
    except Exception:
        log.exception('save task ... exception')
        return False

    log.debug('save task ... successed')
    return True

def has_progress(location, title, filename='progress.cfg'):
    """确定单个任务进度配置文件是否存在
    """

    progress_path = os.path.join(location, title, filename)
    if os.path.isfile(progress_path):
        log.debug('has progress ... yes')
        return True
    else:
        log.debug('has progress ... no | not file: %s', progress_path)
        return False

def check_progress_content(progress):
    """检查任务进度配置内容是否有效
    """

    if not isinstance(progress, int) or progress < -1:
        log.error('check progress content ... error')
        return False

    log.debug('check progress content ... ok')
    return True

def check_progress(progress):
    """对任务进度配置文件做格式校验
    """

    if not isinstance(progress, dict) or not progress:
        log.error('check progress ... error')
        return False

    keys = {'progress'}
    if progress.keys() ^ keys:
        log.error('check progress ... error')
        return False

    if not check_progress_content(progress['progress']):
        log.error('check progress ... error')
        return False

    log.debug('check progress ... ok')
    return True

def load_progress(location, title, filename='progress.cfg'):
    """加载任务进度配置文件

    加载任务进度配置文件，并做格式校验
    """

    task_root_dir = get_task_root_dir(location, title)
    progress_path = os.path.join(task_root_dir, filename)
    if not os.path.isfile(progress_path):
        log.error('load progress ... failed | not file: %s', progress_path)

    progress = {}
    try:
        with open(progress_path, 'r', encoding='utf-8') as f:
            progress = json.load(f)
    except Exception:
        log.exception('load progress ... exception')
        return {}

    if not check_progress(progress):
        log.error('load progress ... failed')
        return {}

    log.debug('load progress ... successed')
    return progress

def save_progress(progress, location, title, filename='progress.cfg'):
    """保存任务进度配置到文件
    """

    if not check_progress(progress):
        log.error('save progress ... failed')
        return False

    task_root_dir = get_task_root_dir(location, title)
    progress_path = os.path.join(task_root_dir, filename)
    try:
        with open(progress_path, 'w', encoding='utf-8') as f:
            json.dump(progress, f, indent=2)
    except Exception:
        log.exception('save progress ... exception')
        return False

    log.debug('save progress ... successed')
    return True

def has_task_config(location, title):
    """确定单个任务配置文件是否存在
    """

    if has_task(location, title) and has_progress(location, title):
        log.info('has task config ... yes')
        return True
    else:
        log.info('has task config ... no')
        return False

def check_task_config(task, progress):
    """对任务配置文件做逻辑校验
    """

    tss = task['tss']
    progress_ = progress['progress']
    if progress_ > len(tss):
        log.error('check task config ... error')
        return False

    log.info('check task config ... ok')
    return True

def get_task_config(location, title):
    """从任务配置文件获取单个任务的配置
    """

    task = load_task(location, title)
    if not task:
        log.error('get task config ... failed')
        return []

    progress = load_progress(location, title)
    if not progress:
        log.error('get task config ... failed')
        return []

    if not check_task_config(task, progress):
        log.error('get task config ... failed')
        return []

    task_config = [
        task,
        progress,
    ]

    log.info('get task config ... successed')
    return task_config

def update_task_config(task_info, only_progress, location, title):
    """将单个任务信息更新到任务配置文件

    由于进度是变动的，而基本任务信息是不变的，所以可以控制是不是只更新进度信息
    """

    if only_progress:
        progress = task_info[1]
        if not save_progress(progress, location, title):
            log.error('update task config ... failed')
            return False
    else:
        task = task_info[0]
        if not save_task(task, location, title):
            log.error('update task config ... failed')
            return False

        progress = task_info[1]
        if not save_progress(progress, location, title):
            log.error('update task config ... failed')
            return False

    log.info('update task config ... successed | only_progress: %s',
             only_progress)
    return True

def is_finished_task(url, location, title):
    """检查是不是已完成的任务

    如果任务的url与断点信息中的url一致，并且断点信息中进度已走完，则认为是已完成的任务
    """

    if not has_task_config(location, title):
        log.info('is finished task ... no | new task')
        return False

    breakpoint_task_config = get_task_config(location, title)
    if not breakpoint_task_config:
        log.warning('is finished task ... no | should not happen')
        return False

    url2 = breakpoint_task_config[0]['url']
    if url2 != url:
        log.info('is finished task ... no | new task')
        return False

    tss = breakpoint_task_config[0]['tss']
    progress_ = breakpoint_task_config[1]['progress']
    if progress_ != len(tss):
        log.info('is finished task ... no | unfinished task')
        return False

    log.info('is finished task ... yes')
    return True

def fetch_file_by_urllib(url, path):
    """从网上取文件并保存到本地

    urllib在多线程环境中使用时，只要让创建的对象在其被创建线程中使用，一般是没有线程安全问题
    的。一些全局设置和全局清理，则要在没有多线程竞争时做，比如在主线程还未创建其他线程时做初
    始设置，在其他线程都结束后再做清理

    urllib自身没有做什么线程安全的措施，其对模块级变量初始化时，是有竞争情况的，不过由于GIL
    确保了对变量、列表和字典的赋值是线程安全的，也就是这些操作不会有中间未决状态，要么就没赋值，
    要么赋好值，所以urllib对模块级变量初始化不会有线程安全问题
    """

    try:
        with urllib.request.urlopen(url) as response:
            with open(path, 'wb') as f:
                shutil.copyfileobj(response, f)
    except Exception:
        log.exception('fetch file by urllib ... exception | url: %s', url)
        return False

    log.debug('fetch file by urllib ... successed | url: %s', url)
    return True

def fetch_file(url, path):
    log.debug('fetch file ... ongoing | args.curl: %s', args.curl)

    if args.curl:
        return util.fetch_file_by_curl(args.curl, url, path)
    else:
        return fetch_file_by_urllib(url, path)

def fetch_data(url):
    """从网上取数据并返回取得数据
    """

    data = b''
    try:
        with urllib.request.urlopen(url) as response:
            data = response.read()
    except Exception:
        log.exception('fetch data ... exception | url: %s', url)
        return b''

    log.debug('fetch data ... successed | url: %s', url)
    return data

def fetch_m3u8_index(url, location, title):
    """获取m3u8的索引文件

    获取m3u8的索引文件，失败会随机间隔地重试指定次数
    """

    task_root_dir = get_task_root_dir(location, title)
    m3u8_index_path = os.path.join(task_root_dir, m3u8_index)
    ok = util.do_with_random_retry(fetch_file, args.retry, args.interval,
                                   url, m3u8_index_path)
    if ok:
        log.info('fetch m3u8 index ... successed')
        return m3u8_index_path
    else:
        log.error('fetch m3u8 index ... failed')
        return ''

def parse_m3u8_index(url, path):
    """分析m3u8索引文件，获取任务配置

    Args:
        url(str): m3u8索引文件的地址
        path(str): m3u8索引文件下载到本地后的存放路径

    Returns:
        如果失败，返回{}

        如果成功，返回如下：

        [
            {
                'url': url,
                'lines': lines,
                'keys': keys,
                'tss': tss,
            },
            {
                'progress': -1,
            },
            {
                'key_data': key_data,
            },
        ]

        url: m3u8索引文件的地址
        lines: 分析后的行的list
        keys: 分析得到的加密信息的list
        tss: 分析后得到的ts片段信息的list
        progress: 已完成进度，-1指示未开始，0指示生成不加密的m3u8索引，正数n指示下载第n个
                  ts片段
        key_data: 占位list，以备之后存放下载得到的密钥数据

    Return type: list

    Raises:
        None
    """

    if not isinstance(path, str) or not path or not os.path.isfile(path):
        log.error('parse m3u8 index ... failed | not file: %s', path)
        return {}

    url_root = util.get_url_root(url)
    if not url_root:
        log.error('parse m3u8 index ... failed | url')
        return {}

    relative_url_head = util.get_relative_url_head(url)
    if not relative_url_head:
        log.error('parse m3u8 index ... failed | url')
        return {}

    try:
        with open(path, 'r', encoding='utf-8') as f:
            sequence = 0
            keys = []
            key_index = 0
            tss = []
            ts_index = 0
            lines = f.readlines()
            for i in range(len(lines)):
                lines[i] = lines[i].strip('\n')

                log.debug(
                    'parse m3u8 index ... ongoing | line: %d, content: %s',
                    i + 1,
                    lines[i])

                yes = is_sequence_line(lines[i])  # 多媒体序列号行
                if yes:
                    sequence = int(yes)
                    continue

                yes = is_key_line(lines[i])  # 加密信息行
                if yes:
                    encryption_info = parse_key_line(lines[i],
                                                     url_root,
                                                     relative_url_head)
                    if not encryption_info:
                        log.error(
                            'parse m3u8 index ... failed | line: %d, '
                            'content: %s',
                            i + 1,
                            lines[i])
                        return {}

                    keys.append(encryption_info)
                    lines[i] = [lines[i], 'k']
                    key_index += 1
                    continue

                yes = is_ts_line(lines[i])  # ts地址行
                if yes:
                    ts_info = parse_ts_line(yes,
                                            url_root,
                                            relative_url_head,
                                            sequence,
                                            key_index - 1)
                    if not ts_info:
                        log.error(
                            'parse m3u8 index ... failed | line: %d, '
                            'content: %s',
                            i + 1,
                            lines[i])
                        return {}

                    tss.append(ts_info)
                    lines[i] = [lines[i], 't', ts_index]
                    ts_index += 1
                    sequence += 1
                    continue
    except Exception:
        log.exception('parse m3u8 index ... exception')
        return {}

    if not check_task_content(url, lines, keys, tss):
        log.error('parse m3u8 index ... failed | invalid content')
        return {}

    task_info = [
        {
            'url': url,
            'lines': lines,
            'keys': keys,
            'tss': tss,
        },
        {
            'progress': -1,
        },
        {
            'key_data': [b'' for i in range(len(keys))],
        },
    ]

    log.info('parse m3u8 index ... successed')
    return task_info

def check_breakpoint_content(lines, breakpoint_lines):
    """检查断点信息的内容是否有效

    两个list长度要一致，对应项如果是str类型，要相等，对应项是list类型，这项的第一个子项要是
    str类型且要相等
    """

    count = len(lines)
    if count != len(breakpoint_lines):
        log.debug('check breakpoint content ... error | unusable')
        return False

    for i in range(count):
        if isinstance(lines[i], str):
            if not isinstance(breakpoint_lines[i], str):
                log.debug('check breakpoint content ... error | unusable')
                return False

            if lines[i] != breakpoint_lines[i]:
                log.debug('check breakpoint content ... error | unusable')
                return False
        elif isinstance(lines[i], list):
            if not isinstance(breakpoint_lines[i], list):
                log.debug('check breakpoint content ... error | unusable')
                return False

            equal = False
            if (lines[i] and breakpoint_lines[i]
                    and isinstance(lines[i][0], str)
                    and isinstance(breakpoint_lines[i][0], str)
                    and lines[i][0] == breakpoint_lines[i][0]):
                equal = True
            if not equal:
                log.debug('check breakpoint content ... error | unusable')
                return False
        else:
            log.warning(
                'check breakpoint content ... error | should not happen')
            return False

    log.debug('check breakpoint content ... ok')
    return True

def check_breakpoint(location, title, task_info):
    """检查断点信息是否可用

    断点信息中的url与任务的url一致，断点信息中进度有效，并且断点信息中的内容也是有效的，则
    认为断点信息有效
    """

    if not has_task_config(location, title):
        log.info('check breakpoint ... error | nonexistent')
        return {}

    breakpoint_task_config = get_task_config(location, title)
    if not breakpoint_task_config:
        log.warning('check breakpoint ... error | should not happen')
        return {}

    url = breakpoint_task_config[0]['url']
    url2 = task_info[0]['url']
    if url != url2:
        log.info('check breakpoint ... error | different task: %s, %s',
                 url,
                 url2)
        return {}

    progress = breakpoint_task_config[1]['progress']
    total_ts = len(task_info[0]['tss'])
    if progress == -1 or progress > total_ts:
        log.info(
            'check breakpoint ... error | useless progress, total: %d, '
            'progress: %d',
            total_ts,
            progress)
        return {}

    if not check_breakpoint_content(task_info[0]['lines'],
                                    breakpoint_task_config[0]['lines']):
        log.info('check breakpoint ... error | different content')
        return {}

    log.info('check breakpoint ... ok')
    return breakpoint_task_config

def update_task_as_needed(task_info, breakpoint_task_config):
    """用断点信息更新任务
    """

    progress = breakpoint_task_config[1]['progress']
    task_info[1]['progress'] = progress
    log.info('update task as needed ... done | progress: %d', progress)

    return task_info

def resume_at_breakpoint(location, title, task_info):
    """根据断点信息恢复任务
    """

    resume = False
    breakpoint_task_config = check_breakpoint(location, title, task_info)
    if breakpoint_task_config:
        resume = True
        task_info = update_task_as_needed(task_info, breakpoint_task_config)

    log.info('resume at breakpoint ... done | resume: %s', resume)
    return resume, task_info

def generate_ts_name(sequence):
    """根据序列号给ts片段命名
    """

    log.debug('generate ts name ... done | sequence: %d', sequence)
    return '{:0>20}'.format(sequence)

def fetch_ts(ts, location, title):
    """获取ts片段文件

    获取ts片段文件，失败会随机间隔地重试指定次数
    """

    url = ts[0]
    sequence = ts[1]

    path = os.path.join(get_m3u8_dir(location, title),
                        generate_ts_name(sequence) + m3u8_extension)
    ok = util.do_with_random_retry(fetch_file, args.retry, args.interval,
                                   url, path)
    if ok:
        log.info('fetch ts ... successed')
        return path
    else:
        log.error('fetch ts ... failed')
        return ''

def fetch_key(url):
    """获取密钥

    获取密钥，失败会随机间隔地重试指定次数
    """

    key = util.do_with_random_retry(fetch_data, args.retry, args.interval, url)
    if key:
        log.info('fetch key ... successed')
    else:
        log.error('fetch key ... failed')

    return key

def decrypt_ts(path, method, key, iv, sequence):
    """解密ts片段文件

    解密后的ts片段文件会替换原来的，key要是bytes类型，iv是str类型
    """

    if (not isinstance(path, str)
            or not path
            or not os.path.isfile(path)
            or not isinstance(method, str)
            or not method
            or not isinstance(key, bytes)
            or not key
            or not isinstance(iv, str)
            or not isinstance(sequence, int)):
        log.error('decrypt ts ... failed | illegal input')
        return False

    log.info('decrypt ts ... ongoing | path: %s, method: %s, key: %s, iv: %s, '
             'sequence: %d'
             % (path, method, key, iv, sequence))

    if method.upper() == 'AES-128':
        try:
            with open(path, 'rb') as f:
                all_data = f.read()

            if iv:
                iv = bytes.fromhex(iv)
            else:
                iv = bytes.fromhex('{:0>32x}'.format(sequence))
            decryptor = AES.new(key, AES.MODE_CBC, iv)
            all_data = decryptor.decrypt(all_data)

            with open(path, 'wb') as f:
                f.write(all_data)
        except Exception:
            log.exception('decrypt ts ... exception')
            return False

        log.info('decrypt ts ... successed')
        return True

    log.error('decrypt ts ... failed')
    return False

def decrypt_ts_if_need(path, ts, keys, key_data):
    """如果ts片段有加密，则进行解密

    如果有加密，则进行解密，密钥获取后存入任务配置，下次则不用再去获取
    """

    total_key = len(keys)
    if total_key != len(key_data):
        log.error('decrypt ts if need ... failed')
        return False

    encrypt_info_index = ts[2]

    if encrypt_info_index < -1 or encrypt_info_index >= total_key:
        log.error('decrypt ts if need ... failed')
        return False

    # -1表示没有加密，所以不用解密
    if encrypt_info_index == -1:
        log.info('decrypt ts if need ... successed | pass')
        return True

    # 加密方法是NONE，这表示没有加密，所以不用解密
    method = keys[encrypt_info_index][0]
    if method.upper() == 'NONE':
        log.info('decrypt ts if need ... successed | pass')
        return True

    key = key_data[encrypt_info_index]
    if not key:
        key_url = keys[encrypt_info_index][1]
        if not key_url:
            log.error('decrypt ts if need ... failed | key url: %s', key_url)
            return False

        key = fetch_key(key_url)
        if not key:
            log.error('decrypt ts if need ... failed | key: %s', key)
            return False

        key_data[encrypt_info_index] = key
        log.info('decrypt ts if need ... ongoing | key: %s', key)

    iv = keys[encrypt_info_index][2]
    sequence = ts[1]
    ok = decrypt_ts(path, method, key, iv, sequence)
    if not ok:
        log.error('decrypt ts if need ... failed')
        return False

    log.info('decrypt ts if need ... successed')
    return True

def generate_m3u8_index(task_info, location, title):
    """生成不加密的m3u8索引文件
    """

    path = os.path.join(get_m3u8_dir(location, title), m3u8_index)

    try:
        with open(path, 'w', encoding='utf-8') as f:
            lines = task_info[0]['lines']
            tss = task_info[0]['tss']
            for line in lines:
                if isinstance(line, str):
                    f.write(line + '\n')
                elif isinstance(line, list):
                    type_ = line[1]
                    if type_ == 't':
                        ts_index = line[2]
                        ts = tss[ts_index]
                        sequence = ts[1]
                        f.write('%s%s\n'
                                % (generate_ts_name(sequence), m3u8_extension))
                else:
                    log.error('generate m3u8 index ... failed')
                    return False
    except Exception:
        log.exception('generate m3u8 index ... exception')
        return False

    log.info('generate m3u8 index ... successed')
    return True

def call_in_thread(func, *args, **kwargs):
    """包装在非主线程中调用的函数，以便在出异常后打印日志信息然后退出
    """

    try:
        yes = func(*args, **kwargs)
    except Exception:
        log.exception('exit from thread ... exception | exception')
        os._exit(1)
    except SystemExit as e:
        log.exception(
            'exit from thread ... exception | systemexit, exit code: %d',
            e.code)
        os._exit(1)

    return yes

def fetch_ts_for_map(ts_info):
    """封装fetch_ts以便使用线程池
    """

    if not isinstance(ts_info, list) and len(ts_info) != 4:
        log.error('fetch ts for pool ... failed ... exit | illegal input')
        sys.exit(1)

    if ts_info[3]:
        log.debug('fetch ts for map ... ignored')
        return (ts_info[3], ts_info[0])
    else:
        ts_info[3] = fetch_ts(ts_info[0], ts_info[1], ts_info[2])
        if ts_info[3]:
            log.debug('fetch ts for map ... successed')
            return (ts_info[3], ts_info[0])
        else:
            log.error('fetch ts for map ... failed')
            return ()

def fetch_ts_in_pool(ts_info):
    return call_in_thread(fetch_ts_for_map, ts_info)

def map2(pool, *args):
    block = pool.map(*args)
    if not all(block):
        return None
    else:
        return block

def fetch_m3u8_data_in_pool(task_info, location, title, parallel):
    """使用线程池获取m3u8的ts片段数据

    如果有加密，则进行解密，每个片段获取成功后更新进度到任务配置文件中
    """

    keys = task_info[0]['keys']
    key_data = task_info[2]['key_data']
    tss = task_info[0]['tss']
    total_ts = len(tss)
    progress = task_info[1]['progress']

    if progress > total_ts or progress < -1:
        log.error(
            'fetch m3u8 data in pool ... failed | total: %d, progress: %d',
            total_ts,
            progress)
        return False

    if progress == total_ts:
        log.warning(
            'fetch m3u8 data in pool ... successed | should not happen',
            total_ts,
            progress)
        return True

    progress += 1  # 进度值表示已完成的步骤，从下一个步骤开始

    # 0号进度值用来指示生成没加密的m3u8索引文件
    if progress == 0:
        if not generate_m3u8_index(task_info, location, title):
            log.error('fetch m3u8 data in pool ... failed | total: %d, now: %d',
                      total_ts,
                      progress)
            return False

        task_info[1]['progress'] = progress
        if not update_task_config(task_info, True, location, title):
            log.error(
                'fetch m3u8 data in pool ... failed | total: %d, now: %d',
                total_ts,
                progress)
            return False

        progress += 1

        log.info('fetch m3u8 data in pool ... ongoing | total: %d, done: %d',
                 total_ts, 0)

    # 1号到total_ts号进度值，分别指示下载第1个到第total_ts个ts片段
    if progress >0 and progress <= total_ts:
        workers = parallel if parallel < max_workers else max_workers

        log.info('fetch m3u8 data in pool ... ongoing | workers: %d', workers)

        with ThreadPool(workers) as pool:
            it = iter(tss[progress - 1 :])
            i = progress - 1

            while True:
                block = tuple(itertools.islice(it, parallel))
                if not block:
                    break

                block2 = [[x, location, title, ''] for x in block]

                log.info('fetch m3u8 data in pool ... ongoing | total: %d, '
                         'now: %d-%d',
                         total_ts,
                         i + 1,
                         i + len(block2))

                block3 = util.do_with_retry(map2, args.retry, pool,
                                            fetch_ts_in_pool, block2)

                if block3 is None:
                    log.error(
                        'fetch m3u8 data in pool ... failed | total: %d, '
                        'now: %d-%d',
                        total_ts,
                        i + 1,
                        i + len(block2))
                    return False

                for x in block3:
                    path = x[0]
                    ts = x[1]

                    if not decrypt_ts_if_need(path, ts, keys, key_data):
                        log.error(
                            'fetch m3u8 data in pool ... failed | total: %d, '
                            'now: %d-%d',
                            total_ts,
                            i + 1,
                            i + len(block2))
                        return False

                task_info[1]['progress'] = i + len(block2)
                if not update_task_config(task_info, True, location, title):
                    log.error(
                        'fetch m3u8 data in pool ... failed | total: %d, '
                        'now: %d-%d',
                        total_ts,
                        i + 1,
                        i + len(block2))
                    return False

                log.info(
                    'fetch m3u8 data in pool ... ongoing | total: %d, '
                    'done: %d-%d',
                    total_ts,
                    i + 1,
                    i + len(block2))

                i += len(block2)

    log.info('fetch m3u8 data in pool ... successed')
    return True

def fetch_m3u8_data(task_info, location, title):
    """获取m3u8的ts片段数据

    如果有加密，则进行解密，每个片段获取成功后更新进度到任务配置文件中
    """

    keys = task_info[0]['keys']
    key_data = task_info[2]['key_data']
    tss = task_info[0]['tss']
    total_ts = len(tss)
    progress = task_info[1]['progress']

    if progress > total_ts or progress < -1:
        log.error(
            'fetch m3u8 data ... failed | total: %d, progress: %d',
            total_ts,
            progress)
        return False

    if progress == total_ts:
        log.warning(
            'fetch m3u8 data ... successed | should not happen',
            total_ts,
            progress)
        return True

    progress += 1  # 进度值表示已完成的步骤，从下一个步骤开始

    # 0号进度值用来指示生成没加密的m3u8索引文件
    if progress == 0:
        if not generate_m3u8_index(task_info, location, title):
            log.error('fetch m3u8 data ... failed | total: %d, now: %d',
                      total_ts,
                      progress)
            return False

        task_info[1]['progress'] = progress
        if not update_task_config(task_info, True, location, title):
            log.error(
                'fetch m3u8 data ... failed | total: %d, now: %d',
                total_ts,
                progress)
            return False

        progress += 1

        log.info('fetch m3u8 data ... ongoing | total: %d, done: %d',
                 total_ts, 0)

    # 1号到total_ts号进度值，分别指示下载第1个到第total_ts个ts片段
    if progress >0 and progress <= total_ts:
        for i in range(progress - 1, total_ts):
            path = fetch_ts(tss[i], location, title)
            if not path:
                log.error(
                    'fetch m3u8 data ... failed | total: %d, now: %d',
                    total_ts,
                    i + 1)
                return False

            if not decrypt_ts_if_need(path, tss[i], keys, key_data):
                log.error(
                    'fetch m3u8 data ... failed | total: %d, now: %d',
                    total_ts,
                    i + 1)
                return False

            task_info[1]['progress'] = i + 1
            if not update_task_config(task_info, True, location, title):
                log.error(
                    'fetch m3u8 data ... failed | total: %d, now: %d',
                    total_ts,
                    i + 1)
                return False

            log.info(
                'fetch m3u8 data ... ongoing | total: %d, done: %d',
                total_ts,
                i + 1)

    log.info('fetch m3u8 data ... successed')
    return True

def fetch_m3u8(url, location, title, parallel):
    """获取m3u8格式视频
    """

    log.info('fetch m3u8 ... ongoing | url: %s, location: %s, title: %s',
             url,
             location,
             title)

    # 看是不是已完成的任务
    yes = is_finished_task(url, location, title)
    util.prompt_true(yes, 'A finished task')
    if yes:
        return

    # 获取m3u8索引
    util.prompt_status('Fetch m3u8 index', 'ongoing')
    m3u8_index_path = fetch_m3u8_index(url, location, title)
    util.prompt_assert(m3u8_index_path, 'Fetch m3u8 index')

    # 分析m3u8索引
    task_info = parse_m3u8_index(url, m3u8_index_path)
    util.prompt_assert(task_info, 'Parse m3u8 index')

    # 断点续传处理
    resume, task_info = resume_at_breakpoint(location, title, task_info)
    util.prompt_true(resume, 'Resume at breakpoint')
    if not resume:
        ok = update_task_config(task_info, False, location, title)
        util.prompt_assert(ok, 'Save task info')

    # 获取m3u8数据
    util.prompt_status('Fetch m3u8 data', 'ongoing')
    if parallel <= 1:
        ok = fetch_m3u8_data(task_info, location, title)
    else:
        ok = fetch_m3u8_data_in_pool(task_info, location, title, parallel)
    util.prompt_assert(ok, 'Fetch m3u8 data')

    log.info('fetch m3u8 ... done')

def check_ffmpeg_output(path, output_path):
    """检查ffmpeg是否成功执行

    暂时还没有办法确认，留待以后发现执行失败的情况的特征，再补充
    """

    log.debug('check ffmpeg output ... ok')
    return True

def convert_m3u8_to_mp4(ffmpeg, location, title):
    """转换m3u8格式到mp4格式
    """

    m3u8_index_path = os.path.join(get_m3u8_dir(location, title), m3u8_index)
    mp4_path = os.path.join(location, title + '.mp4')
    cmd = ('"%s" -n -i "%s" -vcodec copy -acodec copy -absf aac_adtstoasc "%s"'
           % (ffmpeg, m3u8_index_path, mp4_path))
    log.info('convert m3u8 to mp4 ... ongoing | launch: %s', cmd)

    try:
        output_path = m3u8_index_path + '.ffmpegconsole'
        with open(output_path, 'wb') as f:
            exit_code = subprocess.call(cmd, shell=True, stdout=f, stderr=f)
    except Exception:
        log.exception('convert m3u8 to mp4 ... exception')
        return False

    log.info('convert m3u8 to mp4 ... ongoing | exit code: %d', exit_code)
    if exit_code == 0:
        if check_ffmpeg_output(mp4_path, output_path):
            log.info('convert m3u8 to mp4 ... successed')
            return True

    log.error('convert m3u8 to mp4 ... failed')
    return False

def convert(ffmpeg, location, title):
    log.info('convert ... ongoing | ffmpeg: %s, location: %s, title: %s',
             ffmpeg,
             location,
             title)

    util.prompt_status('Convert m3u8 to mp4', 'ongoing')
    ok = convert_m3u8_to_mp4(ffmpeg, location, title)
    util.prompt_assert(ok, 'Convert m3u8 to mp4', False)

    log.info('convert ... done')

def fetch():
    log.info('fetch ... ongoing')

    fetch_m3u8(args.url, args.location, args.title, args.parallel)

    if args.ffmpeg:
        convert(args.ffmpeg, args.location, args.title)

    log.info('fetch ... done')

def is_task_line(line):
    """判断是不是任务内容相关的行

    空白行不是，除了头部空白外以“#”开头行不是，其他的行算是
    """

    if not isinstance(line, str) or not line:
        log.debug('is task line ... no | empty')
        return False

    m = re.search(r'^\s*\#', line)
    if m and m.group():
        log.debug('is task line ... no | comment')
        return False

    m = re.search(r'^\s*$', line)
    if m and m.group():
        log.debug('is task line ... no | space')
        return False

    log.debug('is task line ... yes')
    return True

def get_line_content(line):
    """获取行内容

    去掉头尾空白，剩下内容如果被双引号括起来会去掉双引号
    """

    if not isinstance(line, str) or not line:
        log.error('get line content ... failed | illegal input')
        return ''

    m = re.search(r'^\s*([^\s].*[^\s])\s*$', line)
    if m and m.group() and m.group(1):
        content = m.group(1)
        if content.startswith('"') and content.endswith('"'):
            content = content[1 : len(content) - 1]
            if content:
                log.debug('get line content ... successed')
                return content
            else:
                log.error('get line content ... failed')
                return ''
        else:
            log.debug('get line content ... successed')
            return content

    m = re.search(r'^\s*([^\s])\s*$', line)
    if m and m.group() and m.group(1):
        content = m.group(1)
        log.debug('get line content ... successed')
        return content

    log.error('get line content ... failed | should not happen')
    return ''

def get_task_title(line):
    """获取任务的标题

    会检查是否是支持的标题
    """

    if not isinstance(line, str) or not line:
        log.error('get task title ... failed | illegal input')
        return ''

    title = get_line_content(line)
    if util.is_supported_filename(title):
        log.debug('get task title ... successed')
        return title
    else:
        log.error('get task title ... failed | not supported title')
        return ''

def get_task_url(line):
    """获取任务的URL

    会检查是否是支持的URL
    """

    if not isinstance(line, str) or not line:
        log.error('get task url ... failed | illegal input')
        return ''

    url = get_line_content(line)
    if util.is_supported_url(url):
        log.debug('get task url ... successed')
        return url
    else:
        log.error('get task url ... failed | not supported url')
        return ''

def get_tasks(location, filename='tasks.txt'):
    """从指定位置的特定文件中获取任务清单

    任务清单文件名字必须是“tasks.txt”，其内容是一行标题接着一行URL的格式，会从其中获取URL与
    对应的标题。一个示例如下：

    abc
    https://efg.com/hij/klm.m3u8

    def
    https://hdk.com/slx/dkd.m3u8
    """

    path = os.path.join(location, filename)
    if not os.path.isfile(path):
        log.error('get tasks ... failed | not file: %s', path)
        return []

    try:
        with open(path, 'r', encoding='utf-8') as f:
            lines = [line.strip('\n') for line in f.readlines()
                     if is_task_line(line.strip('\n'))]

            if not lines or len(lines) % 2 != 0:
                log.error('get tasks ... failed | not matched')
                return []

            tasks = []
            unique = []
            for i in range(0, len(lines), 2):
                title = get_task_title(lines[i])
                url = get_task_url(lines[i+1])

                if title and url:
                    if title in unique:
                        log.error('get tasks ... failed | repeated')
                        return []
                    else:
                        tasks.append([url, title])
                        unique.append(title)
                else:
                    log.error('get tasks ... failed | not supported')
                    return []

            log.info('get tasks ... successed')
            return tasks

    except Exception:
        log.exception('get tasks ... exception')
        return []

def fetch_in_sandbox(url, title):
    """执行单个任务的获取操作

    执行单个任务的获取操作，会拦截获取操作导致的任何异常，防止直接退出，以免影响别的任务
    """

    log.info('fetch in sandbox ... ongoing | url: %s, title: %s', url, title)

    try:
        fetch_m3u8(url, args.location, title, args.parallel)

        if args.ffmpeg:
            convert(args.ffmpeg, args.location, title)
    except Exception:
        log.exception('fetch in sandbox ... exception | exception')
        return False
    except SystemExit as e:
        log.exception(
            'fetch in sandbox ... exception | systemexit, exit code: %d',
            e.code)
        return False

    log.info('fetch in sandbox ... done')
    return True

def batch():
    log.info('batch ... ongoing')

    tasks = get_tasks(args.location)
    util.prompt_assert(tasks, 'Get tasks')

    util.prompt('Total %d' % len(tasks))

    failed = []
    for task in tasks:
        util.prompt_status('Fetch %s | %s' % (task[0], task[1]), 'ongoing')
        ok = fetch_in_sandbox(task[0], task[1])
        util.prompt_assert(ok, 'Fetch %s | %s' % (task[0], task[1]), False)

        if not ok:
            failed.append('%s | %s' % (task[0], task[1]))

    util.prompt_batch3(len(tasks), failed)

    log.info('batch ... done')

def arg_type_timeout(str_):
    return util.arg_type_positive2(str_)

def parse_command_line():
    parser = argparse.ArgumentParser(
        prog=__product__,
        description='获取m3u8链接内容到本地以便离线播放')

    parser.add_argument('-v', '--version',
                        action='version',
                        version='%(prog)s ' + __version__)

    subparsers = parser.add_subparsers(
        title='subcommands',
        required=True,
        dest='subcommand',
        description='子命令集，由相应的子命令完成相应的功能')

    common = argparse.ArgumentParser(add_help=False)
    common.add_argument('-p', '--parallel',
                        type=util.arg_type_positive,
                        default=32,
                        help='在获取m3u8的多媒体数据时批量发起请求数，这些请求会提交到'
                             '线程池，然后全部完成后更新进度，较大的值可以提高网络吞吐'
                             '量，默认是32')
    common.add_argument('-t', '--timeout',
                        type=arg_type_timeout,
                        default=30,
                        help='以秒为单位的超时，从网络获取内容没有反应时最多等待时间，'
                             '默认为30秒')
    common.add_argument('-r', '--retry',
                        type=util.arg_type_nonnegative,
                        default=10,
                        help='从网络获取内容出问题时，重试次数，默认为10次')
    common.add_argument('-i', '--interval',
                        type=util.arg_type_positive,
                        nargs='+',
                        default=[2, 4, 8, 16],
                        help='随机重试候选间隔时间系列，单位是秒，默认候选间隔时间系列是'
                             '“2 4 8 16”')
    common.add_argument('-f', '--ffmpeg',
                        type=util.arg_type_ffmpeg,
                        help='ffmpeg命令，若提供，会用其将m3u8转换成mp4，会当成命令行'
                             '命令来执行')
    common.add_argument('-c', '--curl',
                       type=util.arg_type_curl,
                       help='curl命令，若提供，会用其下载，会当成命令行命令来执行')

    # fetch子命令，用于获取单个m3u8链接
    fetch = subparsers.add_parser(
        'fetch',
        help='获取单个m3u8链接',
        parents=[util.arg_log(), util.arg_shutdown(), common])
    fetch.add_argument('url',
                       type=util.arg_type_url,
                       help='m3u8索引文件的链接，包含了m3u8的ts片段信息')
    fetch.add_argument('location',
                       type=util.arg_type_dir,
                       help='本地保存目录')
    fetch.add_argument('title',
                       type=util.arg_type_filename,
                       help='m3u8索引文件关联的内容的标题，以便在本地建立子目录来存储相'
                            '关文件')

    # batch子命令，用于批量获取m3u8链接
    batch = subparsers.add_parser(
        'batch',
        help='批量获取m3u8链接',
        parents=[util.arg_log(), util.arg_shutdown(), common])
    batch.add_argument('location',
                       type=util.arg_type_dir,
                       help='批处理目录，里面需有名为“tasks.txt”的任务清单文件，获取'
                            '到的内容会存储在这个目录下')

    global args
    args = parser.parse_args()

def log_config():
    if hasattr(args, 'log'):
        if hasattr(args, 'log_identifier') and args.log_identifier:
            return (
                '%s_v%s' % (__product__, __version__),
                args.log,
                args.log_identifier
            )
        else:
            return ('%s_v%s' % (__product__, __version__), args.log)

    return None

def init():
    log.info('program version: %s %s', __product__, __version__)
    log.debug('command line: %s', sys.argv)
    log.debug('parse args: %s', args)

    # 获取操作系统密码以便任务完成后自动关机
    if hasattr(args, 'shutdown'):
        global password
        password = util.get_os_password(args.shutdown)

    # 避免目标网站因HTTPS证书问题而导致请求失败
    ssl._create_default_https_context = ssl._create_unverified_context

    # 设置下载超时，这个是全局设置
    socket.setdefaulttimeout(args.timeout)

def schedule_tasks():
    if args.subcommand == 'fetch':
        fetch()
    elif args.subcommand == 'batch':
        batch()

def clean():
    # 自动关机
    if hasattr(args, 'shutdown'):
        util.shutdown(password)
